package com.example.protobuf.config;

import com.google.protobuf.CodedOutputStream;
import com.google.protobuf.Message;
import com.googlecode.protobuf.format.JsonFormat;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.HttpOutputMessage;
import org.springframework.http.MediaType;
import org.springframework.http.converter.AbstractHttpMessageConverter;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.converter.HttpMessageNotReadableException;
import org.springframework.http.converter.HttpMessageNotWritableException;
import org.springframework.http.converter.protobuf.ProtobufHttpMessageConverter;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurationSupport;

import java.io.IOException;
import java.lang.reflect.Method;
import java.nio.charset.Charset;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

import static java.nio.charset.Charset.defaultCharset;

@Configuration
public class MvcConfig extends WebMvcConfigurationSupport {

    @Bean
    RestTemplate restTemplate() {
        RestTemplate rt = new RestTemplate();
        rt.getMessageConverters().add(ownConverter());
        return rt;
    }

    @Override
    public void extendMessageConverters(List<HttpMessageConverter<?>> converters) {
        converters.add(ownConverter());
    }

    private HttpMessageConverter<Message> ownConverter() {
        return new AbstractHttpMessageConverter<Message>() {
            private final ConcurrentHashMap<Class<?>, Method> methodCache = new ConcurrentHashMap<>();
            public final MediaType PROTOBUF = new MediaType("application", "x-protobuf", defaultCharset());

            @Override
            protected boolean supports(Class<?> clazz) {
                return Message.class.isAssignableFrom(clazz);
            }

            @Override
            protected boolean canRead(MediaType mediaType) {
                return true;
            }

            @Override
            protected boolean canWrite(MediaType mediaType) {
                return true;
            }

            @Override
            protected Message readInternal(Class<? extends Message> clazz, HttpInputMessage inputMessage) throws IOException, HttpMessageNotReadableException {
                MediaType contentType = inputMessage.getHeaders().getContentType();
                if (PROTOBUF.isCompatibleWith(contentType)) {
                    try {
                        Message.Builder builder = getMessageBuilder(clazz);
                        builder.mergeFrom(inputMessage.getBody());
                        return builder.build();
                    } catch (Exception ex) {
                        throw new HttpMessageNotReadableException("Could not read Protobuf message: " + ex.getMessage(), ex);
                    }
                }
                throw new HttpMessageNotReadableException("Could not read contentType");
            }

            private Message.Builder getMessageBuilder(Class<? extends Message> clazz) throws Exception {
                Method method = methodCache.get(clazz);
                if (method == null) {
                    method = clazz.getMethod("newBuilder");
                    methodCache.put(clazz, method);
                }
                return (Message.Builder) method.invoke(clazz);
            }


            @Override
            protected void writeInternal(Message message, HttpOutputMessage outputMessage) throws IOException, HttpMessageNotWritableException {
                if (PROTOBUF.isCompatibleWith(outputMessage.getHeaders().getContentType())) {
                    CodedOutputStream codedOutputStream = CodedOutputStream.newInstance(outputMessage.getBody());
                    message.writeTo(codedOutputStream);
                    codedOutputStream.flush();
                }
            }
        };
    }

}
